import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt

feature_params = {'maxCorners': 300, 'qualityLevel': 0.2,
                  'minDistance': 2, 'blockSize': 7}
lk_params = dict(winSize=(15, 15), maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))
file = "C:/Users/18873/Desktop/videos/shibuya.mp4"
cap = cv.VideoCapture(file)

color = (0, 255, 0)
_, first_frame = cap.read()
prev_gray = cv.cvtColor(first_frame, cv.COLOR_BGR2GRAY)
prev = cv.goodFeaturesToTrack(prev_gray, mask=None, **feature_params)
mask = np.zeros_like(first_frame)
# optical_flow = cv.optflow.DualTVL1OpticalFlow_create()


while cap.isOpened():
    _, frame = cap.read()
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    next, status, error = cv.calcOpticalFlowPyrLK(prev_gray, gray, prev, None, **lk_params)

    good_old = prev[status == 1]
    good_new = next[status == 1]

    for i, (new, old) in enumerate(zip(good_new, good_old)):
        a, b = new.ravel()

        c, d = old.ravel()

        mask = cv.line(mask, (a, b), (c, d), color, 2)
        # frame = cv.circle(frame, (a, b), 3, color, -1)
    output = cv.add(frame, mask)
    prev_gray = gray.copy()
    prev = good_new.reshape(-1, 1, 2)
    cv.imshow('sparse optical flow', output)
    # print(output, '\n')

    if cv.waitKey(10) & 0xFF == ord('q'):
        break

cap.release()
cv.destroyAllWindows()
